"""
    该文件的主要作用是分析由相邻参数之间输入图片和隐空间变量之间的关系
    主要包括的内容是：
    （1）控制单维度变化，均等采样图片
    （2）输入网络得到对应的潜空间变量
    （3）求取相邻图片之间的MSE
    （4）可视化二者之间的关系
    （5）保存相关文件
    ## 其主要结构与predict.py相似，可以看做predict文件和analyse的整合，后续可能会与predict、analyse文件整合。
"""

import random
import pandas as pd
import csv
import numpy as np
from utils.base_util import new_fold, get_list_from_dict
from utils.painter import draw_hist
from network.ResNet_model import resnet50
from network.Basic_model import basic_net
from network.metrics import max_distance
import torch
import utils.image_process as imp
from DRR_generator.the_main import Projector
from configs.predict_conf import load_predict_config
from configs.config import global_var
from utils.label_transform import label2real, real2label
import matplotlib.pyplot as plt
from tqdm import tqdm
from utils.painter import draw_line


class Predict:
    def __init__(self, mode, num_classes, model_path, CT_path, saving_drr_path):
        self.mode = mode
        # self.model = resnet50(num_classes=num_classes)
        self.model = basic_net(out_channels=num_classes)
        self.model.load_state_dict(torch.load(model_path))
        # 导入用于DRR生成的对象
        self.data_generator = Projector()
        self.data_generator.load_ct_images(CT_path)
        # 图片保存位置
        self.saving_drr_path = saving_drr_path

    def get_drr(self, pre_alpha, pre_beta, pre_theta, pre_tx, pre_ty, pre_tz):
        new_drr = imp.get_new_drr_numpy(pre_alpha, pre_beta, pre_theta,
                                        pre_tx, pre_ty, pre_tz, self.data_generator)
        return new_drr

    def predict(self, test_img, _tru_r, _tru_t, label):
        if self.mode == 'reg':
            self.model.eval()
            with torch.no_grad():
                test_img = test_img.reshape(1, 1, 128, 128) / 255
                img_tensor = torch.from_numpy(test_img).float()
                out, hidden_vector = self.model(img_tensor)
            pre_r = out[:, 0: 3].detach().cpu().numpy()
            pre_t = out[:, 3: 6].detach().cpu().numpy()
            # 获取预测的DRR
            pre_alpha, pre_beta, pre_theta, pre_tx, pre_ty, pre_tz = label2real(pre_r, pre_t, 1)
            # imp.get_DRR_image(projector=self.data_generator,
            #                   tru_img=test_img,
            #                   tru_r=_tru_r,
            #                   tru_t=_tru_t,
            #                   pre_r=[float(pre_alpha), float(pre_beta), float(pre_theta)],
            #                   pre_t=[float(pre_tx), float(pre_ty), float(pre_tz)],
            #                   i=label,
            #                   saving_path=self.saving_drr_path + "drr_vis_result")
            return np.array([float(pre_alpha), float(pre_beta), float(pre_theta)]), \
                np.array([float(pre_tx), float(pre_ty), float(pre_tz)]), hidden_vector.detach().cpu().numpy()


def dist_visualization(value_name, y1, y2, y3, idx, title, xy_label, path):
    _y1 = get_list_from_dict(y1, idx)
    _y2 = get_list_from_dict(y2, idx)
    _y3 = get_list_from_dict(y3, idx)
    draw_hist(value_name,
              [_y1, _y2, _y3],
              '平移{}:真实值-预测值1-预测值2'.format(title),
              xy_label,
              path)


if __name__ == "__main__":
    new_fold('pre_result', ['noise_result', 'statistical_result'])
    # 导入相应的模型
    pred_conf = load_predict_config(global_var.loading_address)
    predictor = Predict(**pred_conf)
    # 记录运行时的所有真实值，预测值
    tru = pd.DataFrame(columns=['tru_rx', 'tru_ry', 'tru_rz', 'tru_tx', 'tru_ty', 'tru_tz'])
    pre = pd.DataFrame(columns=['pre_rx', 'pre_ry', 'pre_rz', 'pre_tx', 'pre_ty', 'pre_tz'])
    hidden = pd.DataFrame(columns=[('h' + str(i)) for i in range(9248)])
    image = pd.DataFrame(columns=[('im' + str(i)) for i in range(128 * 128)])
    # 每次只对一个维度加噪声
    noise = [0, 0, 0, 0, 0, 0]  # 前三维是角度对应的噪声，后三维是平移对应的噪声
    # 选择采样维度
    img_dict = {"rx": 0, "ry": 1, "rz": 2, "tx": 3, "ty": 4, "tz": 5}
    # 选择采样间隔
    # 角度范围10度，每0.5度采样一次，一共20组
    # 位移范围100mm,每1mm采样一次，一共100组
    sample_interval = {"r": 10, "t": 2}
    img_name = 'tz'  # 实际生成的对比DRR图片的标题
    noise[img_dict[img_name]] = 1
    y1_label = ['绕自身z轴', '绕自身y轴', '绕固定z轴', '沿自身x轴', '沿自身y轴', '沿自身z轴']
    for epoch in tqdm([i / 100 for i in range(-100, 100, sample_interval[img_name[0]])]):
        # print("真值", epoch)
        # 随机生成一组位姿参数, 采用均匀采样的方式
        # 标准正位：（0±15， 270±5， 90±5）
        rx = epoch * 5 * noise[0]
        ry = 270 + epoch * 5 * noise[1]
        rz = 90 + epoch * 5 * noise[2]
        # 位移±50
        tx = epoch * 50 * noise[3]
        ty = epoch * 50 * noise[4]
        tz = epoch * 50 * noise[5]
        # 获取加噪后生成的图片
        test_drr = predictor.get_drr(rx, ry, rz, tx, ty, tz).reshape(128, 128, 1)
        # 通过网络获取预测的输出
        pre_r, pre_t, hidden_pre = predictor.predict(test_drr, [rx, ry, rz], [tx, ty, tz], epoch)
        # 根据网络的输出投影出对应的DRR
        pre_drr = predictor.get_drr(pre_r[0], pre_r[1], pre_r[2],
                                    pre_t[0], pre_t[1], pre_t[2]).reshape(128, 128, 1)
        # 数据记录
        if img_name[0] == 'r':
            row = round((epoch + 1) * 10)
        elif img_name[0] == 't':
            row = round((epoch + 1) * 50)
        else:
            print("无效前缀！！！，请输入以下前缀：")
            print('[rx, ry, rz, tx, ty, tz]')
            break
        # print("图片序号", row)
        tru.loc[row] = [rx, ry, rz, tx, ty, tz]
        pre.loc[row] = [pre_r[0], pre_r[1], pre_r[2], pre_t[0], pre_t[1], pre_t[2]]
        hidden.loc[row] = hidden_pre.ravel()
        image.loc[row] = test_drr.ravel()
        # 可视化
        this_tru = np.array([rx, ry, rz, tx, ty, tz]).ravel()
        this_pre = np.array([pre_r[0], pre_r[1], pre_r[2], pre_t[0], pre_t[1], pre_t[2]]).ravel()
        plt.subplot(1, 2, 1)
        plt.title('实际值'
                  + "\n" +
                  '角度' + str(np.around(this_tru[0:3], 2))
                  + "\n" +
                  '平移' + str(np.around(this_tru[3:], 2)))
        plt.imshow(test_drr, cmap='gray')
        plt.subplot(1, 2, 2)
        plt.title('预测值'
                  + "\n" +
                  '角度' + str(np.around(this_pre[0:3], 2))
                  + "\n" +
                  '平移' + str(np.around(this_pre[3:], 2)))
        plt.imshow(pre_drr, cmap='gray')
        plt.suptitle("{0}:验证{1}加噪的精度".format(row, img_name))
        plt.savefig(predictor.saving_drr_path + "noise_result/第{0}组.png".format(row))
        # plt.show()
    # 写入csv文件
    tru.to_csv(predictor.saving_drr_path + "statistical_result/tru.csv")
    pre.to_csv(predictor.saving_drr_path + "statistical_result/pre.csv")
    hidden.to_csv(predictor.saving_drr_path + "statistical_result/hidden.csv")
    image.to_csv(predictor.saving_drr_path + "statistical_result/test_image.csv")
    # 获取真值和预测值之间的差
    tru = tru['tru_{0}'.format(img_name)]
    pre = pre['pre_{0}'.format(img_name)]
    pre_error = np.round(abs(pre - tru), 2)
    # print(pre_error.shape)
    # 记录每张图片的id
    test_id = [(i + 1) / 100 for i in range(pre_error.shape[0])]
    # 计算相邻参数图片和隐空间之间的距离
    image_margin = np.linalg.norm(image - image.shift(1), axis=1)[1:] / 255
    hidden_margin = np.linalg.norm(hidden - hidden.shift(1), axis=1)[1:].astype(np.float64)
    # print(type(image_margin[0]))
    # print(type(hidden_margin[0]))
    # 为了作图美观，把横坐标的中间值归到0。仅对标准正位有效
    if img_name == 'ry':
        tru = tru - 270
    elif img_name == 'rz':
        tru = tru - 90
    draw_line(tru, [np.round(image_margin, 2), np.round(hidden_margin, 2), pre_error, test_id],
              ['image_margin', 'hidden_margin', 'pre_error', 'id'],
              predictor.saving_drr_path + "statistical_result/img_HD_{}.html".format(img_name), img_name,
              mode='one_x_axis', title=y1_label[noise.index(1)])
